﻿using System.Collections.Concurrent;

namespace MhAdmin.Web.Core;

public class ConnectionMiddleware
{
    private readonly RequestDelegate _next;
    private readonly ConcurrentDictionary<string, SemaphoreSlim> _concurrentConnections;
    private readonly int _maxConcurrentRequests;

    public ConnectionMiddleware(RequestDelegate next, int maxConcurrentRequests)
    {
        _next = next;
        _concurrentConnections = new ConcurrentDictionary<string, SemaphoreSlim>();
        _maxConcurrentRequests = maxConcurrentRequests;
    }

    public async Task Invoke(HttpContext context)
    {
        string requestPath = context.Request.Path;
        string method = context.Request.Method;

        // 检查是否有此IP地址的现有信号量
        var semaphore = _concurrentConnections.GetOrAdd(requestPath, new SemaphoreSlim(_maxConcurrentRequests));

        try
        {
            if (method != "OPTIONS" && !semaphore.Wait(0))
            {
                // 等待信号量
                context.Response.StatusCode = 209;
                return;
            }

            // 调用管道中的下一个中间件
            await _next(context);
        }
        finally
        {
            // 释放信号量
            semaphore.Release();
        }
    }
}
